#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
星火X1 WebSocket服务
基于讯飞星火X1 WebSocket API
"""

import websocket
import datetime
import hashlib
import base64
import hmac
import json
import logging
import ssl
import threading
import time
from urllib.parse import urlparse, urlencode
from wsgiref.handlers import format_date_time
from datetime import datetime
from time import mktime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SparkWebSocketParam:
    """星火X1 WebSocket参数类"""
    
    def __init__(self, app_id, api_key, api_secret, spark_url):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.host = urlparse(spark_url).netloc
        self.path = urlparse(spark_url).path
        self.spark_url = spark_url
    
    def create_url(self):
        """生成WebSocket连接URL"""
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        
        # 拼接签名字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"
        
        # 进行hmac-sha256加密
        signature_sha = hmac.new(
            self.api_secret.encode('utf-8'), 
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()
        
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
        
        authorization_origin = f'api_key="{self.api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        
        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        
        # 拼接鉴权参数，生成url
        url = self.spark_url + '?' + urlencode(v)
        return url


class SparkWebSocketService:
    """星火X1 WebSocket服务"""
    
    def __init__(self, app_id, api_key, api_secret, spark_url="wss://spark-api.xf-yun.com/v1/x1"):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.spark_url = spark_url
        self.answer = ""
        self.is_first_content = False
        self.is_completed = False
        self.error_message = None
        self.ws = None
        
        logger.info("星火X1 WebSocket服务初始化完成")
    
    def _on_message(self, ws, message):
        """处理WebSocket消息"""
        try:
            data = json.loads(message)
            code = data['header']['code']
            
            if code != 0:
                error_msg = f"请求错误: {code}, {data}"
                logger.error(error_msg)
                self.error_message = error_msg
                self.is_completed = True
                ws.close()
                return
            
            choices = data["payload"]["choices"]
            status = choices["status"]
            text = choices['text'][0]
            
            # 处理推理内容
            if 'reasoning_content' in text and text['reasoning_content']:
                reasoning_content = text["reasoning_content"]
                logger.info(f"推理内容: {reasoning_content}")
                self.is_first_content = True
            
            # 处理回答内容
            if 'content' in text and text['content']:
                content = text["content"]
                if self.is_first_content:
                    logger.info("*******************以上为思维链内容，模型回复内容如下********************")
                logger.info(f"回答内容: {content}")
                self.is_first_content = False
                self.answer += content
            
            # 检查是否完成
            if status == 2:
                logger.info("回答完成")
                self.is_completed = True
                ws.close()
                
        except Exception as e:
            logger.error(f"处理消息异常: {e}")
            self.error_message = f"处理消息异常: {e}"
            self.is_completed = True
    
    def _on_error(self, ws, error):
        """处理WebSocket错误"""
        logger.error(f"WebSocket错误: {error}")
        self.error_message = f"WebSocket错误: {error}"
        self.is_completed = True
    
    def _on_close(self, ws, *args):
        """处理WebSocket关闭"""
        logger.info("WebSocket连接已关闭")
        self.is_completed = True
    
    def _on_open(self, ws):
        """处理WebSocket连接建立"""
        def run():
            data = json.dumps(self.request_data)
            logger.info("开始发送请求数据")
            ws.send(data)
        
        threading.Thread(target=run).start()
    
    def _generate_params(self, question, domain="x1", temperature=1.2, max_tokens=32768, 
                        enable_web_search=False, search_mode="normal"):
        """
        生成请求参数
        
        Args:
            question: 对话历史列表
            domain: 模型版本
            temperature: 温度参数
            max_tokens: 最大token数
            enable_web_search: 是否启用网络搜索
            search_mode: 搜索模式
            
        Returns:
            dict: 请求参数
        """
        data = {
            "header": {
                "app_id": self.app_id,
                "uid": "1234",
            },
            "parameter": {
                "chat": {
                    "domain": domain,
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                    "presence_penalty": 1,
                    "frequency_penalty": 0.02,
                    "top_k": 5
                }
            },
            "payload": {
                "message": {
                    "text": question
                }
            }
        }
        
        # 添加网络搜索工具
        if enable_web_search:
            data["parameter"]["chat"]["tools"] = [
                {
                    "type": "web_search",
                    "web_search": {
                        "enable": True,
                        "search_mode": search_mode
                    }
                }
            ]
        
        return data
    
    def chat_completion(self, question, domain="x1", temperature=1.2, max_tokens=32768,
                       enable_web_search=False, search_mode="normal", stream=True):
        """
        聊天完成
        
        Args:
            question: 对话历史列表
            domain: 模型版本
            temperature: 温度参数
            max_tokens: 最大token数
            enable_web_search: 是否启用网络搜索
            search_mode: 搜索模式
            stream: 是否流式输出
            
        Returns:
            str: AI回答
        """
        logger.info("开始聊天完成")
        
        # 重置状态
        self.answer = ""
        self.is_first_content = False
        self.is_completed = False
        self.error_message = None
        
        # 生成请求参数
        self.request_data = self._generate_params(
            question=question,
            domain=domain,
            temperature=temperature,
            max_tokens=max_tokens,
            enable_web_search=enable_web_search,
            search_mode=search_mode
        )
        
        # 创建WebSocket参数
        ws_param = SparkWebSocketParam(
            app_id=self.app_id,
            api_key=self.api_key,
            api_secret=self.api_secret,
            spark_url=self.spark_url
        )
        
        # 生成WebSocket URL
        ws_url = ws_param.create_url()
        
        try:
            # 创建WebSocket连接
            websocket.enableTrace(False)
            self.ws = websocket.WebSocketApp(
                ws_url,
                on_message=self._on_message,
                on_error=self._on_error,
                on_close=self._on_close,
                on_open=self._on_open
            )
            
            # 运行WebSocket连接
            self.ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
            
            # 检查是否有错误
            if self.error_message:
                raise Exception(self.error_message)
            
            logger.info(f"聊天完成，回答长度: {len(self.answer)} 字符")
            return self.answer
            
        except Exception as e:
            logger.error(f"聊天完成失败: {e}")
            raise Exception(f"聊天完成失败: {e}")


class ConversationManager:
    """对话历史管理器"""
    
    def __init__(self, max_length=8000):
        self.conversation_history = []
        self.max_length = max_length
    
    def add_message(self, role, content):
        """添加消息到对话历史"""
        message = {"role": role, "content": content}
        self.conversation_history.append(message)
        
        # 检查长度并删除旧消息
        while self._get_total_length() > self.max_length:
            if len(self.conversation_history) > 1:
                del self.conversation_history[0]
            else:
                break
        
        return self.conversation_history.copy()
    
    def _get_total_length(self):
        """获取对话总长度"""
        total_length = 0
        for message in self.conversation_history:
            total_length += len(message["content"])
        return total_length
    
    def get_history(self):
        """获取对话历史"""
        return self.conversation_history.copy()
    
    def clear_history(self):
        """清空对话历史"""
        self.conversation_history = []


def main():
    """测试函数"""
    print("=== 星火X1 WebSocket服务测试 ===")
    
    # 从配置文件加载配置
    with open('config.json', 'r', encoding='utf-8') as f:
        config = json.load(f)
    
    # 创建星火X1 WebSocket服务
    llm_config = config['llm']
    spark_service = SparkWebSocketService(
        app_id=config['app_id'],
        api_key=llm_config['api_key'].replace('Bearer ', ''),
        api_secret=config['access_key_secret']
    )
    
    # 创建对话管理器
    conversation_manager = ConversationManager()
    
    # 测试对话
    test_questions = [
        "你好，请介绍一下你自己",
        "请解释一下什么是机器学习",
        "推荐一本编程书籍"
    ]
    
    for i, question in enumerate(test_questions, 1):
        print(f"\n--- 测试 {i} ---")
        print(f"用户: {question}")
        
        try:
            # 添加用户消息
            conversation_manager.add_message("user", question)
            
            # 获取回答
            answer = spark_service.chat_completion(
                question=conversation_manager.get_history(),
                enable_web_search=False
            )
            
            print(f"助手: {answer}")
            
            # 添加助手回答
            conversation_manager.add_message("assistant", answer)
            
        except Exception as e:
            print(f"对话失败: {e}")
    
    # 显示对话历史
    print("\n--- 对话历史 ---")
    for i, msg in enumerate(conversation_manager.get_history()):
        role = "用户" if msg["role"] == "user" else "助手"
        content = msg["content"][:50] + "..." if len(msg["content"]) > 50 else msg["content"]
        print(f"{i+1}. {role}: {content}")


if __name__ == "__main__":
    main()
